import os

import fsspec
from pyarrow.fs import FSSpecHandler, PyFileSystem

# 使用 fsspec 创建一个 S3 文件系统对象
s3_fs = fsspec.filesystem(
    's3',
    key='nvx1',
    secret='nvx1',
    client_kwargs={'endpoint_url': 'http://192.168.16.114:38333'})

# 将 S3 文件系统包装为 PyArrow 文件系统
pyarrow_fs = PyFileSystem(FSSpecHandler(s3_fs))

file_path = 's3://photon/mlflow/1/2257a8ab2bea4968811bdbf2daca28a3/artifacts/'
# 获取文件信息
file_info = pyarrow_fs.get_file_info(file_path)
print(file_info)
print(dir(file_info.type))
print(file_info.type.name)
print(file_info.type.value)

# # 读取文件内容
# with pyarrow_fs.open_input_stream(file_path) as f:
#     content = f.read()
#     print(content)

if file_info.type.name.lower() == 'directory':
    entries = s3_fs.ls(file_path)
    for s3_file_path in entries:
        entry = pyarrow_fs.get_file_info(s3_file_path)
        print(f'Name: {entry.path}, Type: {entry.type}, Size: {entry.size}')
        local_file_path = os.path.basename(entry.path)
        with pyarrow_fs.open_input_file(s3_file_path) as source_file, open(
                local_file_path, 'wb') as local_file:
            while True:
                chunk = source_file.read(1024 * 1024)  # 每次读取 1MB
                if not chunk:
                    break
                local_file.write(chunk)
else:
    print(f'{file_path} is not a directory.')

# 指定本地文件路径和目标 S3 文件路径
local_file_path = './input.json'
s3_file_path = 's3://photon/mlflow/1/2257a8ab2bea4968811bdbf2daca28a3/artifacts1/12/12/12/input.json'

# 打开文件流进行读取和写入
with open(local_file_path, 'rb') as local_file, pyarrow_fs.open_output_stream(
        s3_file_path) as s3_file:
    while True:
        chunk = local_file.read(1024 * 1024)  # 每次读取 1MB
        if not chunk:
            break
        s3_file.write(chunk)
